# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import annotations

import abc
import re
import shlex
from dataclasses import dataclass
from typing import TYPE_CHECKING, ClassVar, Final, Iterable, Set

from typing_extensions import override

import crossbench.path as pth
from crossbench import exception
from crossbench.parse import ObjectParser
from crossbench.probes.probe import Probe, ProbeConfigParser, ProbeContext
from crossbench.probes.result_location import ResultLocation

if TYPE_CHECKING:
  import datetime as dt

  from crossbench.probes.results import ProbeResult
  from crossbench.runner.run import Run


class DownloadsProbe(Probe):
  """
  Probe that captures downloads from websites and allows loading tests to wait
  for a download to complete.
  """
  NAME: ClassVar = "downloads"
  RESULT_LOCATION: ClassVar = ResultLocation.BROWSER

  CHROME_OS_DOWNLOADS_DIR = pth.AnyPath("/home/chronos/user/MyFiles/Downloads")

  @classmethod
  def config_parser(cls) -> ProbeConfigParser:
    parser = super().config_parser()
    parser.add_argument(
        "clear_downloads",
        aliases=("clear",),
        type=ObjectParser.bool,
        default=False,
        help="Delete all files in the download folder before every run.")
    parser.add_argument(
        "save_downloads",
        aliases=("save",),
        type=ObjectParser.bool,
        default=True,
        help=("Copy all files downloaded during a test to the test results"
              " folder."))
    return parser

  def __init__(self,
               clear_downloads: bool = False,
               save_downloads: bool = False) -> None:
    super().__init__()
    self._clear_downloads: bool = clear_downloads
    self._save_downlaods: bool = save_downloads

  @override
  def create_context(self, run: Run) -> DownloadsProbeContext:
    if run.browser_platform.is_android:
      return AndroidWebDriverDownloadsProbeContext(self, run)

    if run.browser_platform.is_chromeos:
      return FileWatchDownloadsProbeContext(self, run,
                                            self.CHROME_OS_DOWNLOADS_DIR)
    raise NotImplementedError(
        f"Probe({self}): Unsupported browser: {run.browser}")

  @property
  def clear_downloads(self) -> bool:
    return self._clear_downloads

  @property
  def save_downloads(self) -> bool:
    return self._save_downlaods


class DownloadsProbeContext(ProbeContext[DownloadsProbe]):

  def __init__(self, probe: DownloadsProbe, run: Run) -> None:
    super().__init__(probe, run)

  def get_default_result_path(self) -> pth.AnyPath:
    downloads_dir = super().get_default_result_path()
    self.browser_platform.mkdir(downloads_dir)
    return downloads_dir

  @override
  def invoke(self, info_stack: exception.TInfoStack, timeout: dt.timedelta,
             **kwargs) -> None:
    del info_stack
    self._wait_for_download(timeout, **kwargs)

  def _wait_for_download(self, timeout: dt.timedelta, pattern: str | re.Pattern,
                         **kwargs) -> None:
    self.expect_no_extra_kwargs(kwargs)

    if not isinstance(pattern, re.Pattern):
      pattern = re.compile(pattern)

    wait_range = self.run.wait_range(min_interval=0.2, timeout=timeout)
    for _ in wait_range.wait_with_backoff():
      if self.download_complete(pattern):
        return

  @abc.abstractmethod
  def download_complete(self, pattern: re.Pattern) -> bool:
    pass


class FileWatchDownloadsProbeContext(DownloadsProbeContext):

  def __init__(self, probe: DownloadsProbe, run: Run,
               downloads_dir: pth.AnyPath) -> None:
    super().__init__(probe, run)
    self._downloads_dir: pth.AnyPath = downloads_dir
    self._existing_downloads: Set[pth.AnyPath] = set()
    self._results: list[pth.AnyPath] = []

  def downloads(self, include_pending: bool = True) -> Iterable[pth.AnyPath]:
    downloads = self.browser_platform.iterdir(self._downloads_dir)
    if include_pending:
      return downloads
    return [file for file in downloads if file.suffix != ".crdownload"]

  def start(self) -> None:
    if self.probe.clear_downloads:
      for file in self.downloads():
        self.browser_platform.rm(file)
      self._existing_downloads = set()
    else:
      self._existing_downloads = set(self.downloads())

  def stop(self) -> None:
    if not self.probe.save_downloads:
      return
    for file in self.downloads():
      if file in self._existing_downloads:
        continue
      to_path = self.result_path / file.name
      self.browser_platform.copy_file(file, to_path)
      self._results.append(to_path)

  def teardown(self) -> ProbeResult:
    return self.browser_result(file=self._results)

  @override
  def download_complete(self, pattern: re.Pattern) -> bool:
    return any(pattern.search(file.name) for file in self.downloads())


@dataclass(frozen=True)
class AndroidDownload:
  id: str
  display_name: str


class AndroidWebDriverDownloadsProbeContext(DownloadsProbeContext):
  CONTENT_QUERY_RE: Final[re.Pattern] = re.compile(
      r"Row: \d+ _display_name=(.*), _id=(\d+)")
  CONTENT_QUERY_NO_RESULTS: Final = "No result found."

  def __init__(self, probe: DownloadsProbe, run: Run) -> None:
    super().__init__(probe, run)
    self._existing_downloads: Set[AndroidDownload] = set()
    self._user_id: str = str(self.browser_platform.user_id())
    self._results: list[pth.AnyPath] = []

  def downloads(self,
                include_pending: bool = True) -> Iterable[AndroidDownload]:
    result: list[AndroidDownload] = []
    args = [
        "content", "query", "--user", self._user_id, "--uri",
        "content://media/external/downloads", "--where", "is_download=1",
        "--projection", "_display_name:_id"
    ]
    if not include_pending:
      args.append("--where")
      args.append("is_pending=0")
    rows = self.browser_platform.sh_stdout(*args)
    if rows.strip() == self.CONTENT_QUERY_NO_RESULTS:
      return result
    for row in rows.splitlines():
      if match := self.CONTENT_QUERY_RE.match(row):
        result.append(
            AndroidDownload(display_name=match.group(1), id=match.group(2)))
      else:
        raise RuntimeError(
            f"Android downloads content query unexpect result row: {row}")

    return result

  def delete(self, download: AndroidDownload) -> None:
    self.browser_platform.sh("content", "delete", "--user", self._user_id,
                             "--uri", "content://media/external/downloads",
                             "--where", f"_id='{download.id}'")

  def start(self) -> None:
    if self.probe.clear_downloads:
      for download in self.downloads():
        self.delete(download)
    self._existing_downloads = set(self.downloads())

  def stop(self) -> None:
    if not self.probe.save_downloads:
      return
    for download in self.downloads():
      if download in self._existing_downloads:
        continue
      to_path = self.result_path / download.display_name
      read_downloads_cmd = (
          "content",
          "read",
          "--user",
          self._user_id,
          "--uri",
          f"content://media/external/downloads/{download.id}",
      )
      cmd = (
          shlex.join(read_downloads_cmd) + ">" +
          shlex.quote(self.browser_platform.path(to_path).as_posix()))
      # We need shell=True since we pipe to a file.
      self.browser_platform.sh(cmd, shell=True)  # noqa: S604
      self._results.append(to_path)

  def teardown(self) -> ProbeResult:
    return self.browser_result(file=self._results)

  @override
  def download_complete(self, pattern: re.Pattern) -> bool:
    downloads = self.downloads(include_pending=False)
    return any(pattern.search(download.display_name) for download in downloads)
